#include "mesh_renderer.h"
#include <cassert>
#include <cmath>

namespace vkrender
{
    MeshRenderer::MeshRenderer(
        const std::shared_ptr<VulkanDeviceContext>& context,
        const std::shared_ptr<VulkanSwapChain>& swapchain)
    {
        this->context = context;
        this->swapchain = swapchain.get();
        ubo.proj = glm::mat4(1.0f);
        ubo.view = glm::mat4(1.0f);
        ubo.model = glm::mat4(1.0f);

        // load resources
        const std::string texture_path = "../resources/models/viking_room.png";
        texture = std::make_shared<VulkanTexture>(context, texture_path);
        sampler = std::make_shared<VulkanTextureSampler>(context);

        // createCommandBuffers();
        createUniformBuffers();
        createDescriptorSetLayout();
        createPipelineLayout();
        createDescriptorPool();
        createDescriptorSets();

        updateDescriptorSets();
        adjustForSwapchain(swapchain);
    }

    MeshRenderer::~MeshRenderer()
    {
    }

    void MeshRenderer::createDescriptorSetLayout()
    {
        // create descriptor set layout
        VkDescriptorSetLayoutBinding samplerLayoutBinding = {};
        samplerLayoutBinding.binding = 0;
        samplerLayoutBinding.descriptorCount = 1;
        samplerLayoutBinding.descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
        samplerLayoutBinding.pImmutableSamplers = nullptr;
        samplerLayoutBinding.stageFlags = VK_SHADER_STAGE_FRAGMENT_BIT;

        VkDescriptorSetLayoutBinding uboLayoutBinding = {};
        uboLayoutBinding.binding = 1;
        uboLayoutBinding.descriptorCount = 1;
        uboLayoutBinding.descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
        uboLayoutBinding.pImmutableSamplers = nullptr;
        uboLayoutBinding.stageFlags = VK_SHADER_STAGE_VERTEX_BIT;

        std::array<VkDescriptorSetLayoutBinding, 2> bindings = { samplerLayoutBinding, uboLayoutBinding };
        VkDescriptorSetLayoutCreateInfo layoutInfo = {};
        layoutInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_LAYOUT_CREATE_INFO;
        layoutInfo.bindingCount = static_cast<uint32_t>(bindings.size());
        layoutInfo.pBindings = bindings.data();

        vkCreateDescriptorSetLayout(context->device, &layoutInfo, nullptr, &descriptorSetLayout);
    }

    void MeshRenderer::createPipelineLayout()
    {
        VkPipelineLayoutCreateInfo pipelineLayoutInfo{};
        pipelineLayoutInfo.sType = VK_STRUCTURE_TYPE_PIPELINE_LAYOUT_CREATE_INFO;
        pipelineLayoutInfo.setLayoutCount = 1;
        pipelineLayoutInfo.pSetLayouts = &descriptorSetLayout;
        pipelineLayoutInfo.pushConstantRangeCount = 0;
        pipelineLayoutInfo.pPushConstantRanges = nullptr;

        if (vkCreatePipelineLayout(context->device, &pipelineLayoutInfo, nullptr, &pipelineLayout) != VK_SUCCESS)
        {
            throw std::runtime_error("failed to create pipeline layout!");
        }
    }

    void MeshRenderer::createPipeline()
    {
        assert(swapchain != nullptr && "Cannot create pipeline before swap chain");
        assert(pipelineLayout != nullptr && "Cannot create pipeline before pipeline layout");

        PipelineConfigInfo pipelineConfig{};
        // fill in graphics pipeline configurations.
        VulkanGraphicsPipeline::defaultPipelineConfigInfo(pipelineConfig);
        pipelineConfig.renderPass = swapchain->renderpass;
        pipelineConfig.pipelineLayout = pipelineLayout;

        auto bindingDescriptions = MeshVertex::getBindingDescriptions();
        auto attributeDescriptions = MeshVertex::getAttributeDescriptions();

        graphicsPipeline = std::make_shared<VulkanGraphicsPipeline>(
            context,
            "../resources/shaders/mesh.vert.spv",
            "../resources/shaders/mesh.frag.spv",
            bindingDescriptions, attributeDescriptions,
            pipelineConfig);
    }
#if 0
    void MeshRenderer::createCommandBuffers()
    {
        commandBuffers.resize(swapchain->imageCount());

        VkCommandBufferAllocateInfo allocInfo{};
        allocInfo.sType = VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO;
        allocInfo.level = VK_COMMAND_BUFFER_LEVEL_PRIMARY;
        allocInfo.commandPool = context->commandPool;
        allocInfo.commandBufferCount = static_cast<uint32_t>(commandBuffers.size());

        if (vkAllocateCommandBuffers(context->device, &allocInfo, commandBuffers.data()) !=
            VK_SUCCESS)
        {
            throw std::runtime_error("failed to allocate command buffers!");
        }
    }

    void MeshRenderer::freeCommandBuffers()
    {
        if (commandBuffers.size() == 0)
            return;

        vkFreeCommandBuffers(
            context->device,
            context->commandPool,
            static_cast<uint32_t>(commandBuffers.size()),
            commandBuffers.data());
        commandBuffers.clear();
    }
#endif

    void MeshRenderer::adjustForSwapchain(const std::shared_ptr<VulkanSwapChain>& swapchain)
    {
        this->swapchain = swapchain.get();
#if 0
        if (swapchain->imageCount() != commandBuffers.size()) {
            freeCommandBuffers();
            createCommandBuffers();
        }
#endif
        createPipeline();
    }

    void MeshRenderer::bind(uint32_t imageIndex, VkCommandBuffer cmd,
        const glm::mat4& proj, const glm::mat4& view)
    {
        ubo.proj = proj;
        ubo.view = view;
        graphicsPipeline->bind(cmd);

        vkCmdBindDescriptorSets(cmd, VK_PIPELINE_BIND_POINT_GRAPHICS,
            pipelineLayout, 0, 1, &descriptorSets[imageIndex], 0, nullptr);
    }

    void MeshRenderer::draw(uint32_t imageIndex, VkCommandBuffer cmd,
        const std::shared_ptr<VulkanModel>& model,
        const glm::mat4& model_matrix)
    {
        ubo.model = model_matrix;
        model->bind(cmd);
        model->draw(cmd);
        updateUniformBuffer(imageIndex);
    }

#if 0
    VkCommandBuffer MeshRenderer::endDraw(uint32_t imageIndex)
    {
        
        vkCmdEndRenderPass(commandBuffers[imageIndex]);
        if (vkEndCommandBuffer(commandBuffers[imageIndex]) != VK_SUCCESS)
        {
            throw std::runtime_error("failed to record command buffer!");
        }
        return commandBuffers[imageIndex];
    }
#endif

    void MeshRenderer::createDescriptorPool()
    {
        uint32_t num_frames = swapchain->imageCount();
        // create descriptor pool
        std::array<VkDescriptorPoolSize, 2> poolSizes = {};
        poolSizes[0].type = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
        poolSizes[0].descriptorCount = static_cast<uint32_t>(num_frames);
        poolSizes[1].type = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
        poolSizes[1].descriptorCount = static_cast<uint32_t>(num_frames);
        VkDescriptorPoolCreateInfo poolInfo = {};
        poolInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_POOL_CREATE_INFO;
        poolInfo.poolSizeCount = static_cast<uint32_t>(poolSizes.size());
        poolInfo.pPoolSizes = poolSizes.data();
        poolInfo.maxSets = static_cast<uint32_t>(num_frames);
        vkCreateDescriptorPool(context->device, &poolInfo, nullptr, &descriptorPool);
    }

    void MeshRenderer::createDescriptorSets()
    {
        uint32_t num_frames = swapchain->imageCount();

        // create descriptor sets
        std::vector<VkDescriptorSetLayout> layouts(num_frames, descriptorSetLayout);
        VkDescriptorSetAllocateInfo allocInfo = {};
        allocInfo.sType = VK_STRUCTURE_TYPE_DESCRIPTOR_SET_ALLOCATE_INFO;
        allocInfo.descriptorPool = descriptorPool;
        allocInfo.descriptorSetCount = static_cast<uint32_t>(num_frames);
        allocInfo.pSetLayouts = layouts.data();

        descriptorSets.resize(num_frames);
        vkAllocateDescriptorSets(context->device, &allocInfo, descriptorSets.data());
    }

    void MeshRenderer::updateDescriptorSets()
    {
        uint32_t num_frames = swapchain->imageCount();

        for (uint32_t i = 0; i < num_frames; ++i)
        {
            VkDescriptorImageInfo imageInfo = {};
            imageInfo.imageLayout = VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL;
            imageInfo.imageView = texture->getImageView();
            imageInfo.sampler = sampler->getSampler();

            VkDescriptorBufferInfo bufferInfo = {};
            bufferInfo.buffer = uniformBuffers[i].buffer;
            bufferInfo.offset = 0;
            bufferInfo.range = sizeof(MVPUBO);

            std::array<VkWriteDescriptorSet, 2> descriptorWrite = {};
            descriptorWrite[0].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
            descriptorWrite[0].dstSet = descriptorSets[i];
            descriptorWrite[0].dstBinding = 0;
            descriptorWrite[0].dstArrayElement = 0;
            descriptorWrite[0].descriptorType = VK_DESCRIPTOR_TYPE_COMBINED_IMAGE_SAMPLER;
            descriptorWrite[0].descriptorCount = 1;
            descriptorWrite[0].pImageInfo = &imageInfo;

            descriptorWrite[1].sType = VK_STRUCTURE_TYPE_WRITE_DESCRIPTOR_SET;
            descriptorWrite[1].dstSet = descriptorSets[i];
            descriptorWrite[1].dstBinding = 1;
            descriptorWrite[1].dstArrayElement = 0;
            descriptorWrite[1].descriptorType = VK_DESCRIPTOR_TYPE_UNIFORM_BUFFER;
            descriptorWrite[1].descriptorCount = 1;
            descriptorWrite[1].pBufferInfo = &bufferInfo;
            vkUpdateDescriptorSets(context->device, static_cast<uint32_t>(descriptorWrite.size()),
                descriptorWrite.data(), 0, nullptr);
        }
    }

    void MeshRenderer::createUniformBuffers()
    {
        uint32_t num_frames = swapchain->imageCount();
        VkDeviceSize buffer_size = sizeof(MVPUBO);
        for (uint32_t i = 0; i < num_frames; ++i)
        {
            uniformBuffers.emplace_back(context, buffer_size);
        }
    }

    void MeshRenderer::updateUniformBuffer(uint32_t imageIndex)
    {
        uniformBuffers[imageIndex].updateBufferData(ubo);
    }
}
